# Initialize CPML boundary condition    
from define_problem_space_parameters import boundary,dx,dy,dz
from initialize_boundary_conditions import is_cpml_xn,n_cpml_xn,is_cpml_xp,n_cpml_xp,is_cpml_yn,\
    n_cpml_yn,is_cpml_yp,n_cpml_yp,is_cpml_zn,n_cpml_zn,is_cpml_zp,n_cpml_zp
from initialize_fdtd_parameters_and_arrays import mu_0,eps_0,dt
from calculate_domain_size import ny,nzp1,nyp1,nz,nxp1,nx
from initialize_updating_coefficients import Ceyhz,Cezhy,Chzey,Chyez,Cezhx,Cexhz,Chzex,Chxez,Cexhy,\
    Cexhy,Ceyhx,Chxey,Chyex
import numpy as np
p_order = boundary["cpml_order"] # order of the polynomial distribution
sigma_ratio = boundary["cpml_sigma_factor"]
kappa_max = boundary["cpml_kappa_max"]
alpha_min = boundary["cpml_alpha_min"]
alpha_max = boundary["cpml_alpha_max"]

# Initialize cpml for xn region
if is_cpml_xn:
    # define one-dimensional temporary cpml parameter arrays 
    sigma_max = sigma_ratio  * (p_order+1)/(150*np.pi*dx)
    ncells = n_cpml_xn
    rho_e = (np.arange(ncells, 0, -1) - 0.75) / ncells
    rho_m = (np.arange(ncells, 0, -1) - 0.25) / ncells
    sigma_pex_xn = sigma_max * rho_e**p_order
    sigma_pmx_xn = sigma_max * rho_m**p_order
    sigma_pmx_xn = (mu_0/eps_0) * sigma_pmx_xn
    kappa_ex_xn = 1 + (kappa_max - 1) * rho_e**p_order
    kappa_mx_xn = 1 + (kappa_max - 1) * rho_m**p_order
    alpha_ex_xn = alpha_min + (alpha_max - alpha_min) * (1-rho_e)
    alpha_mx_xn = alpha_min + (alpha_max - alpha_min) * (1-rho_m)
    alpha_mx_xn = (mu_0/eps_0) * alpha_mx_xn
    
    # define one-dimensional cpml parameter arrays 
    cpml_b_ex_xn = np.exp((-dt/eps_0)*((sigma_pex_xn/kappa_ex_xn)+ alpha_ex_xn)) 
    cpml_a_ex_xn = (1/dx)*(cpml_b_ex_xn-1.0)* sigma_pex_xn/(kappa_ex_xn*(sigma_pex_xn+kappa_ex_xn*alpha_ex_xn))
    cpml_b_mx_xn = np.exp((-dt/mu_0)*((sigma_pmx_xn/kappa_mx_xn)+ alpha_mx_xn)) 
    cpml_a_mx_xn = (1/dx)*(cpml_b_mx_xn-1.0) * sigma_pmx_xn/(kappa_mx_xn*(sigma_pmx_xn+kappa_mx_xn*alpha_mx_xn))

    # Create and initialize 2D cpml convolution parameters 
    Psi_eyx_xn = np.zeros((ncells,ny,nzp1)) 
    Psi_ezx_xn = np.zeros((ncells,nyp1,nz)) 
    Psi_hyx_xn = np.zeros((ncells,nyp1,nz)) 
    Psi_hzx_xn = np.zeros((ncells,ny,nzp1)) 

    # Create and initialize 2D cpml convolution coefficients 
    # Notice that Ey(1,:,:) and Ez(1,:,:) are not updated by cmpl 
    CPsi_eyx_xn = Ceyhz[1:ncells+1, :, :] * dx  # 注意Python是0-based索引
    CPsi_ezx_xn = Cezhy[1:ncells+1, :, :] * dx
    CPsi_hyx_xn = Chyez[0:ncells, :, :] * dx    # 相当于MATLAB的1:ncells
    CPsi_hzx_xn = Chzey[0:ncells, :, :] * dx
    
    # Adjust FDTD coefficients in the CPML region 
    # Notice that Ey(1,:,:) and Ez(1,:,:) are not updated by cmpl 
    for i in range(ncells):
        Ceyhz[i+1, :, :] /= kappa_ex_xn[i]    # 注意Python是0-based索引
        Cezhy[i+1, :, :] /= kappa_ex_xn[i]
        Chyez[i, :, :] /= kappa_mx_xn[i]
        Chzey[i, :, :] /= kappa_mx_xn[i]
        
    # Delete temporary arrays. These arrays will not be used any more.
    del sigma_pex_xn,sigma_pmx_xn
    del kappa_ex_xn ,kappa_mx_xn
    del alpha_ex_xn ,alpha_mx_xn

# Initialize cpml for xp region
if is_cpml_xp:
    # define one-dimensional temporary cpml parameter arrays 
    sigma_max = sigma_ratio  * (p_order+1)/(150*np.pi*dx)
    ncells = n_cpml_xp
    rho_e = (np.arange(1, ncells + 1) - 0.75) / ncells
    rho_m = (np.arange(1, ncells + 1) - 0.25) / ncells
    sigma_pex_xp = sigma_max * rho_e**p_order
    sigma_pmx_xp = sigma_max * rho_m**p_order
    sigma_pmx_xp = (mu_0/eps_0) * sigma_pmx_xp
    kappa_ex_xp = 1 + (kappa_max - 1) * rho_e**p_order
    kappa_mx_xp = 1 + (kappa_max - 1) * rho_m**p_order
    alpha_ex_xp = alpha_min + (alpha_max - alpha_min) * (1-rho_e)
    alpha_mx_xp = alpha_min + (alpha_max - alpha_min) * (1-rho_m)
    alpha_mx_xp = (mu_0/eps_0) * alpha_mx_xp
    
    # define one-dimensional cpml parameter arrays 
    cpml_b_ex_xp = np.exp((-dt/eps_0) \
        *((sigma_pex_xp/kappa_ex_xp)+ alpha_ex_xp)) 
    cpml_a_ex_xp = (1/dx)*(cpml_b_ex_xp-1.0)* sigma_pex_xp \
        /(kappa_ex_xp*(sigma_pex_xp+kappa_ex_xp*alpha_ex_xp))
    cpml_b_mx_xp = np.exp((-dt/mu_0) \
        *((sigma_pmx_xp/kappa_mx_xp)+ alpha_mx_xp)) 
    cpml_a_mx_xp = (1/dx)*(cpml_b_mx_xp-1.0) * sigma_pmx_xp \
        /(kappa_mx_xp*(sigma_pmx_xp+kappa_mx_xp*alpha_mx_xp))

    # Create and initialize 2D cpml convolution parameters 
    Psi_eyx_xp = np.zeros((ncells,ny,nzp1)) 
    Psi_ezx_xp = np.zeros((ncells,nyp1,nz)) 
    Psi_hyx_xp = np.zeros((ncells,nyp1,nz)) 
    Psi_hzx_xp = np.zeros((ncells,ny,nzp1)) 

    # Create and initialize 2D cpml convolution coefficients 
    # Notice that Ey(nxp1,:,:) and Ez(nxp1,:,:) are not updated by cmpl 
    CPsi_eyx_xp = Ceyhz[nxp1-ncells-1:nx,:,:]*dx
    CPsi_ezx_xp = Cezhy[nxp1-ncells-1:nx,:,:]*dx
    CPsi_hyx_xp = Chyez[nxp1-ncells-1:nx,:,:]*dx
    CPsi_hzx_xp = Chzey[nxp1-ncells-1:nx,:,:]*dx
    
    # Adjust FDTD coefficients in the CPML region 
    # Notice that Ey(nxp1,:,:) and Ez(nxp1,:,:) are not updated by cmpl 
    for i in range(ncells):
        idx = nx - ncells + i  # 计算当前索引
        Ceyhz[idx, :, :] /= kappa_ex_xp[i]
        Cezhy[idx, :, :] /= kappa_ex_xp[i]
        Chyez[idx, :, :] /= kappa_mx_xp[i]
        Chzey[idx, :, :] /= kappa_mx_xp[i]
    
    # Delete temporary arrays. These arrays will not be used any more.
    del sigma_pex_xp ,sigma_pmx_xp
    del kappa_ex_xp ,kappa_mx_xp
    del alpha_ex_xp ,alpha_mx_xp



# Initialize cpml for yn region
if is_cpml_yn:

    # define one-dimensional temporary cpml parameter arrays 
    sigma_max = sigma_ratio  * (p_order+1)/(150*np.pi*dy)
    ncells = n_cpml_yn
    rho_e = ( np.arange(ncells, 0, -1) - 0.75) / ncells
    rho_m = ( np.arange(ncells, 0, -1) - 0.25) / ncells
    sigma_pey_yn = sigma_max * rho_e**p_order
    sigma_pmy_yn = sigma_max * rho_m**p_order
    sigma_pmy_yn = (mu_0/eps_0) * sigma_pmy_yn
    kappa_ey_yn = 1 + (kappa_max - 1) * rho_e**p_order
    kappa_my_yn = 1 + (kappa_max - 1) * rho_m**p_order
    alpha_ey_yn = alpha_min + (alpha_max - alpha_min) * (1-rho_e)
    alpha_my_yn = alpha_min + (alpha_max - alpha_min) * (1-rho_m)
    alpha_my_yn = (mu_0/eps_0) * alpha_my_yn
    
    # define one-dimensional cpml parameter arrays 
    cpml_b_ey_yn = np.exp((-dt/eps_0) \
        *((sigma_pey_yn/kappa_ey_yn)+ alpha_ey_yn)) 
    cpml_a_ey_yn = (1/dy)*(cpml_b_ey_yn-1.0)* sigma_pey_yn \
        /(kappa_ey_yn*(sigma_pey_yn+kappa_ey_yn*alpha_ey_yn))
    cpml_b_my_yn = np.exp((-dt/mu_0) \
        *((sigma_pmy_yn/kappa_my_yn)+ alpha_my_yn)) 
    cpml_a_my_yn = (1/dy)*(cpml_b_my_yn-1.0) * sigma_pmy_yn \
        /(kappa_my_yn*(sigma_pmy_yn+kappa_my_yn*alpha_my_yn))

    # Create and initialize 2D cpml convolution parameters 
    Psi_ezy_yn = np.zeros((nxp1,ncells,nz)) 
    Psi_exy_yn = np.zeros((nx,ncells,nzp1)) 
    Psi_hzy_yn = np.zeros((nx,ncells,nzp1)) 
    Psi_hxy_yn = np.zeros((nxp1,ncells,nz)) 

    # Create and initialize 2D cpml convolution coefficients 
    # Notice that Ez(:,1,:) and Ex(:,1,:) are not updated by cmpl 
    CPsi_ezy_yn = Cezhx[:,1:ncells+1,:]*dy
    CPsi_exy_yn = Cexhz[:,1:ncells+1,:]*dy
    CPsi_hzy_yn = Chzex[:,0:ncells,:]*dy
    CPsi_hxy_yn = Chxez[:,0:ncells,:]*dy
    
    # Adjust FDTD coefficients in the CPML region 
    # Notice that Ez(:,1,:) and Ex(:,1,:) are not updated by cmpl 
    for i in range(ncells):
        Cezhx[:, i + 1, :] /= kappa_ey_yn[i]  # 对应 MATLAB 的 i+1
        Cexhz[:, i + 1, :] /= kappa_ey_yn[i]
        Chzex[:, i, :]     /= kappa_my_yn[i]  # 对应 MATLAB 的 i
        Chxez[:, i, :]     /= kappa_my_yn[i]
    # Delete temporary arrays. These arrays will not be used any more.
    del sigma_pey_yn ,sigma_pmy_yn
    del kappa_ey_yn ,kappa_my_yn
    del alpha_ey_yn ,alpha_my_yn

# Initialize cpml for yp region
if is_cpml_yp:

    # define one-dimensional temporary cpml parameter arrays 
    sigma_max = sigma_ratio  * (p_order+1)/(150*np.pi*dy)
    ncells = n_cpml_yp
    rho_e = (np.arange(1, ncells + 1) - 0.75) / ncells
    rho_m = (np.arange(1, ncells + 1) - 0.25) / ncells
    sigma_pey_yp = sigma_max * rho_e**p_order
    sigma_pmy_yp = sigma_max * rho_m**p_order
    sigma_pmy_yp = (mu_0/eps_0) * sigma_pmy_yp
    kappa_ey_yp = 1 + (kappa_max - 1) * rho_e**p_order
    kappa_my_yp = 1 + (kappa_max - 1) * rho_m**p_order
    alpha_ey_yp = alpha_min + (alpha_max - alpha_min) * (1-rho_e)
    alpha_my_yp = alpha_min + (alpha_max - alpha_min) * (1-rho_m)
    alpha_my_yp = (mu_0/eps_0) * alpha_my_yp
    
    # define one-dimensional cpml parameter arrays 
    cpml_b_ey_yp = np.exp((-dt/eps_0) \
        *((sigma_pey_yp/kappa_ey_yp)+ alpha_ey_yp)) 
    cpml_a_ey_yp = (1/dy)*(cpml_b_ey_yp-1.0)* sigma_pey_yp \
        /(kappa_ey_yp*(sigma_pey_yp+kappa_ey_yp*alpha_ey_yp))
    cpml_b_my_yp = np.exp((-dt/mu_0) \
        *((sigma_pmy_yp/kappa_my_yp)+ alpha_my_yp)) 
    cpml_a_my_yp = (1/dy)*(cpml_b_my_yp-1.0) * sigma_pmy_yp \
        /(kappa_my_yp*(sigma_pmy_yp+kappa_my_yp*alpha_my_yp))

    # Create and initialize 2D cpml convolution parameters 
    Psi_ezy_yp = np.zeros((nxp1,ncells,nz)) 
    Psi_exy_yp = np.zeros((nx,ncells,nzp1)) 
    Psi_hzy_yp = np.zeros((nx,ncells,nzp1)) 
    Psi_hxy_yp = np.zeros((nxp1,ncells,nz)) 

    # Create and initialize 2D cpml convolution coefficients 
    # Notice that Ez(:,nyp1,:) and Ex(:,nyp1,:) are not updated by cmpl 
    CPsi_ezy_yp = Cezhx[:,nyp1-ncells-1:ny,:]*dy
    CPsi_exy_yp = Cexhz[:,nyp1-ncells-1:ny,:]*dy
    CPsi_hzy_yp = Chzex[:,nyp1-ncells-1:ny,:]*dy
    CPsi_hxy_yp = Chxez[:,nyp1-ncells-1:ny,:]*dy
    
    # Adjust FDTD coefficients in the CPML region 
    # Notice that Ez(:,nyp1,:) and Ex(:,nyp1,:) are not updated by cmpl 
    for i in range(ncells):
        # 计算当前索引位置（从ny-ncells开始）
        # 对各个数组进行除法操作
        Cezhx[:, ny - ncells + i, :] /= kappa_ey_yp[i]
        Cexhz[:, ny - ncells + i, :] /= kappa_ey_yp[i]
        Chzex[:, ny - ncells + i, :] /= kappa_my_yp[i]
        Chxez[:, ny - ncells + i, :] /= kappa_my_yp[i]
    
    # Delete temporary arrays. These arrays will not be used any more.
    del sigma_pey_yp ,sigma_pmy_yp
    del kappa_ey_yp ,kappa_my_yp
    del alpha_ey_yp ,alpha_my_yp

# Initialize cpml for zn region
if is_cpml_zn:

    # define one-dimensional temporary cpml parameter arrays 
    sigma_max = sigma_ratio  * (p_order+1)/(150*np.pi*dz)
    ncells = n_cpml_zn
    rho_e = (np.arange(ncells, 0, -1) - 0.75) / ncells
    rho_m = (np.arange(ncells, 0, -1) - 0.25) / ncells
    sigma_pez_zn = sigma_max * rho_e**p_order
    sigma_pmz_zn = sigma_max * rho_m**p_order
    sigma_pmz_zn = (mu_0/eps_0) * sigma_pmz_zn
    kappa_ez_zn = 1 + (kappa_max - 1) * rho_e**p_order
    kappa_mz_zn = 1 + (kappa_max - 1) * rho_m**p_order
    alpha_ez_zn = alpha_min + (alpha_max - alpha_min) * (1-rho_e)
    alpha_mz_zn = alpha_min + (alpha_max - alpha_min) * (1-rho_m)
    alpha_mz_zn = (mu_0/eps_0) * alpha_mz_zn
    
    # define one-dimensional cpml parameter arrays 
    cpml_b_ez_zn = np.exp((-dt/eps_0) \
        *((sigma_pez_zn/kappa_ez_zn)+ alpha_ez_zn)) 
    cpml_a_ez_zn = (1/dz)*(cpml_b_ez_zn-1.0)* sigma_pez_zn \
        /(kappa_ez_zn*(sigma_pez_zn+kappa_ez_zn*alpha_ez_zn))
    cpml_b_mz_zn = np.exp((-dt/mu_0) \
        *((sigma_pmz_zn/kappa_mz_zn)+ alpha_mz_zn)) 
    cpml_a_mz_zn = (1/dz)*(cpml_b_mz_zn-1.0) * sigma_pmz_zn \
        /(kappa_mz_zn*(sigma_pmz_zn+kappa_mz_zn*alpha_mz_zn))

    # Create and initialize 2D cpml convolution parameters 
    Psi_exz_zn = np.zeros((nx,nyp1,ncells)) 
    Psi_eyz_zn = np.zeros((nxp1,ny,ncells)) 
    Psi_hxz_zn = np.zeros((nxp1,ny,ncells)) 
    Psi_hyz_zn = np.zeros((nx,nyp1,ncells)) 

    # Create and initialize 2D cpml convolution coefficients 
    # Notice that Ex(:,:,1) and Ey(:,:,1) are not updated by cmpl 
    CPsi_exz_zn = Cexhy[:,:,1:ncells+1]*dz
    CPsi_eyz_zn = Ceyhx[:,:,1:ncells+1]*dz
    CPsi_hxz_zn = Chxey[:,:,0:ncells]*dz
    CPsi_hyz_zn = Chyex[:,:,0:ncells]*dz
    
    # Adjust FDTD coefficients in the CPML region 
    # Notice that Ex(:,:,1) and Ey(:,:,1) are not updated by cmpl 
    for i in np.arange(ncells):                                       
        Cexhy[:,:,i+1] /=kappa_ez_zn[i]
        Ceyhx[:,:,i+1] /=kappa_ez_zn[i]
        Chxey[:,:,i] /=kappa_mz_zn[i]
        Chyex[:,:,i] /=kappa_mz_zn[i]
    
    # Delete temporary arrays. These arrays will not be used any more.
    del sigma_pez_zn, sigma_pmz_zn
    del kappa_ez_zn, kappa_mz_zn
    del alpha_ez_zn, alpha_mz_zn



# Initialize cpml for zp region
if is_cpml_zp:

    # define one-dimensional temporary cpml parameter arrays 
    sigma_max = sigma_ratio  * (p_order+1)/(150*np.pi*dz)
    ncells = n_cpml_zp
    rho_e = (np.arange(ncells)+0.25)/ncells
    rho_m = (np.arange(ncells)+0.75)/ncells
    sigma_pez_zp = sigma_max * rho_e**p_order
    sigma_pmz_zp = sigma_max * rho_m**p_order
    sigma_pmz_zp = (mu_0/eps_0) * sigma_pmz_zp
    kappa_ez_zp = 1 + (kappa_max - 1) * rho_e**p_order
    kappa_mz_zp = 1 + (kappa_max - 1) * rho_m**p_order
    alpha_ez_zp = alpha_min + (alpha_max - alpha_min) * (1-rho_e)
    alpha_mz_zp = alpha_min + (alpha_max - alpha_min) * (1-rho_m)
    alpha_mz_zp = (mu_0/eps_0) * alpha_mz_zp
    
    # define one-dimensional cpml parameter arrays 
    cpml_b_ez_zp = np.exp((-dt/eps_0) \
        *((sigma_pez_zp/kappa_ez_zp)+ alpha_ez_zp)) 
    cpml_a_ez_zp = (1/dz)*(cpml_b_ez_zp-1.0)* sigma_pez_zp \
        /(kappa_ez_zp*(sigma_pez_zp+kappa_ez_zp*alpha_ez_zp))
    cpml_b_mz_zp = np.exp((-dt/mu_0) \
        *((sigma_pmz_zp/kappa_mz_zp)+ alpha_mz_zp)) 
    cpml_a_mz_zp = (1/dz)*(cpml_b_mz_zp-1.0) * sigma_pmz_zp \
        /(kappa_mz_zp*(sigma_pmz_zp+kappa_mz_zp*alpha_mz_zp))

    # Create and initialize two-dimensional cpml convolution parameters 
    Psi_exz_zp = np.zeros((nx,nyp1,ncells)) 
    Psi_eyz_zp = np.zeros((nxp1,ny,ncells)) 
    Psi_hxz_zp = np.zeros((nxp1,ny,ncells)) 
    Psi_hyz_zp = np.zeros((nx,nyp1,ncells)) 

    # Create and initialize 2D cpml convolution coefficients 
    # Notice that Ex(:,:,nzp1) and Ey(:,:,nzp1) are not updated by cmpl 
    CPsi_exz_zp = Cexhy[:,:,nzp1-ncells-1:nz]*dz
    CPsi_eyz_zp = Ceyhx[:,:,nzp1-ncells-1:nz]*dz
    CPsi_hxz_zp = Chxey[:,:,nzp1-ncells-1:nz]*dz
    CPsi_hyz_zp = Chyex[:,:,nzp1-ncells-1:nz]*dz
    
    # Adjust FDTD coefficients in the CPML region 
    # Notice that Ex(:,:,nzp1) and Ey(:,:,nzp1) are not updated by cmpl 
    for i in np.arange(ncells):
        Cexhy[:,:,nz-ncells+i] /= kappa_ez_zp[i]
        Ceyhx[:,:,nz-ncells+i] /= kappa_ez_zp[i]
        Chxey[:,:,nz-ncells+i] /= kappa_mz_zp[i]
        Chyex[:,:,nz-ncells+i] /= kappa_mz_zp[i]
    
    # Delete temporary arrays. These arrays will not be used any more.
    del sigma_pez_zp, sigma_pmz_zp
    del kappa_ez_zp, kappa_mz_zp
    del alpha_ez_zp, alpha_mz_zp
